# Copyright 2024 Marimo. All rights reserved.
from __future__ import annotations

from typing import Optional, Union

from marimo._config.config import CopilotMode
from marimo._server.models.completion import (
    AiCompletionContext,
    Language,
    SchemaTable,
    VariableContext,
)
from marimo._types.ids import SessionId

FIM_PREFIX_TAG = "<|fim_prefix|>"
FIM_SUFFIX_TAG = "<|fim_suffix|>"
FIM_MIDDLE_TAG = "<|fim_middle|>"

LANGUAGES: list[Language] = ["python", "sql", "markdown"]
language_rules: dict[Language, list[str]] = {
    "python": [
        "For matplotlib: use plt.gca() as the last expression instead of plt.show().",
        "For plotly: return the figure object directly.",
        "For altair: return the chart object directly. Add tooltips where appropriate. You can pass polars dataframes directly to altair (e.g., alt.Chart(df)).",
        "Include proper labels, titles, and color schemes.",
        "Make visualizations interactive where appropriate.",
        "If an import already exists, do not import it again.",
        "If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.",
    ],
    "markdown": [],
    "sql": [
        "The SQL must use duckdb syntax.",
    ],
}


language_rules_multiple_cells: dict[Language, list[str]] = {
    "sql": [
        'SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines. You should always write queries inline as the code snippet above, do not use variables to store queries.',
        "This will automatically display the result in the UI. You do not need to return the dataframe in the cell.",
        "The SQL must use the syntax of the database engine specified in the `engine` variable. If no engine, then use duckdb syntax.",
    ]
}


def _format_schema_info(tables: Optional[list[SchemaTable]]) -> str:
    """Helper to format schema information from context"""
    if not tables:
        return ""

    schema_info = "\n\n## Available schema:\n"
    for schema in tables:
        schema_info += f"- Table: {schema.name}\n"
        for col in schema.columns:
            schema_info += f"  - Column: {col.name}\n"
            schema_info += f"    - Type: {col.type}\n"
            if col.sample_values:
                samples = ", ".join(f"{v}" for v in col.sample_values)
                schema_info += f"    - Sample values: {samples}\n"
    return schema_info


def _format_plain_text(plain_text: str) -> str:
    if not plain_text.strip():
        return ""
    return f"If the prompt mentions @kind://name, use the following context to help you answer the question:\n\n{plain_text}"


def _format_variables(
    variables: Optional[list[Union[VariableContext, str]]],
) -> str:
    if not variables:
        return ""

    variable_info = "\n\n## Available variables from other cells:\n"
    for variable in variables:
        if isinstance(variable, VariableContext):
            if _is_private_variable := variable.name.startswith("_"):
                continue
            variable_info += f"- variable: `{variable.name}`\n"
            variable_info += f"  - value_type: {variable.value_type}\n"
            variable_info += f"  - value_preview: {variable.preview_value}\n"
        else:
            if _is_private_variable := variable.startswith("_"):
                continue
            variable_info += f"- variable: `{variable}`"

    return variable_info


def _rules(rules: list[str]) -> str:
    """Format a list of rules into a numbered string."""
    return "\n".join(f"{i + 1}. {rule}" for i, rule in enumerate(rules))


def get_refactor_or_insert_notebook_cell_system_prompt(
    *,
    language: Language,
    is_insert: bool,
    support_multiple_cells: bool,
    custom_rules: Optional[str],
    cell_code: Optional[str],
    selected_text: Optional[str],
    other_cell_codes: Optional[str],
    context: Optional[AiCompletionContext],
) -> str:
    if cell_code:
        system_prompt = f"Here's a {language} document from a Python notebook that I'm going to ask you to make an edit to.\n\n"
    elif support_multiple_cells:
        system_prompt = (
            "You are an AI assistant integrated into the marimo notebook code editor.\n"
            "Your goal is to create new cells in the notebook.\n"
            "You can create multiple cells with different languages. Each cell should be wrapped in backticks.\n"
            "The user may reference additional context in the form @kind://name. You can use this context to help you with the current task.\n"
            "You can reference variables from other cells, but you cannot redefine a variable if it already exists.\n"
            "Immediately start with the following format. Do NOT comment on the code, just output the code itself: \n\n"
            "```python\n{PYTHON_CODE}\n```\n\n"
            '```sql\ndf_name = mo.sql(f"""{SQL_QUERY}""")\n```\n\n'
            '```markdown\nmo.md(f"""{MARKDOWN_CONTENT}""")\n```\n\n'
            "You can have multiple cells of any type. Each cell is wrapped in backticks with the appropriate language identifier.\n"
            "Create clear variable names if they will be used in other cells. Do not prefix with underscore.\n"
            "Separate logic into multiple cells to keep the code organized and readable."
        )
    else:
        system_prompt = (
            "You are an AI assistant integrated into the marimo notebook code editor.\n"
            "You goal is to create a new cell in the notebook.\n"
            f"Your output must be valid {language} code.\n"
            "The user may reference additional context in the form @kind://name. You can use this context to help you with the current task.\n"
            "You can reference variables from other cells, but you cannot redefine a variable if it already exists.\n"
            "Immediately start with the following format. Do NOT comment on the code, just output the code itself: \n\n"
            "```\n{CELL_CODE}\n```"
        )

    # When we are modifying or inserting into an existing cell, we need to
    if cell_code:
        # Assertions, otherwise the system prompt will be wrong
        if is_insert:
            assert "<insert_here>" in cell_code
            system_prompt += "The point you'll need to insert at is marked with <insert_here></insert_here>.\n"
        else:
            # If no selected text and <rewrite_this> is not present, we need to
            # wrap the whole code in <rewrite_this> tags
            if "<rewrite_this>" not in cell_code and not selected_text:
                cell_code = f"<rewrite_this>\n{cell_code}\n</rewrite_this>"

            assert "<rewrite_this>" in cell_code
            system_prompt += "The section you'll need to rewrite is marked with <rewrite_this></rewrite_this> tags.\n"

        if cell_code:
            system_prompt += "\n\n" + _tag("document", cell_code) + "\n\n"

        if is_insert:
            system_prompt += (
                "You can't replace the content, your answer will be inserted in place of the "
                "<insert_here></insert_here> tags. Don't include the insert_here tags in your output.\n"
                "Match the indentation in the original file in the inserted content, "
                "don't include any indentation on blank lines.\n"
                "Immediately start with the following format. Do NOT comment on the code, just output the code itself:\n\n"
                "```\n{INSERTED_CODE}\n```"
            )
        else:
            system_prompt += (
                "Only make changes that are necessary to fulfill the prompt, leave everything else as-is. "
                "All surrounding content will be preserved.\n"
                "Start at the indentation level in the original file in the rewritten content. "
                "Don't stop until you've rewritten the entire section, even if you have no more changes to make, "
                "always write out the whole section with no unnecessary elisions.\n"
                "Immediately start with the following format. Do NOT comment on the code, just output the code itself:\n\n"
                "```\n{REWRITTEN_CODE}\n```"
            )

    if selected_text:
        system_prompt += "\n\nAnd here's the section to rewrite based on that prompt again for reference:\n\n"
        system_prompt += _tag("rewrite_this", selected_text)

    if support_multiple_cells:
        # Add all language rules for multi-cell scenarios
        for lang in LANGUAGES:
            language_rule = language_rules_multiple_cells.get(
                lang, language_rules.get(lang, [])
            )
            if language_rule:
                system_prompt += (
                    f"\n\n## Rules for {lang}:\n{_rules(language_rule)}"
                )
    elif language in language_rules and language_rules[language]:
        system_prompt += (
            f"\n\n## Rules for {language}\n{_rules(language_rules[language])}"
        )

    if custom_rules and custom_rules.strip():
        system_prompt += f"\n\n## Additional rules:\n{custom_rules}"

    if context:
        system_prompt += _format_plain_text(context.plain_text)
        system_prompt += _format_variables(context.variables)
        system_prompt += _format_schema_info(context.schema)

    if other_cell_codes:
        system_prompt += "\n\n" + _tag(
            "code_from_other_cells", other_cell_codes
        )

    if support_multiple_cells:
        system_prompt += "\n\nAgain, just output code wrapped in cells. Each cell is wrapped in backticks with the appropriate language identifier (python, sql, markdown)."
    else:
        system_prompt += f"\n\nAgain, just output the code itself and make sure to return the code as just {language}."

    return system_prompt


def get_inline_system_prompt(*, language: Language) -> str:
    return (
        f"You are a {language} code completion assistant. "
        f"Complete the missing code between the prefix and suffix while maintaining proper syntax, style, and functionality."
        f"Only output the code that goes after the {FIM_SUFFIX_TAG} part. Do not add any explanation or markdown."
    )


def _get_mode_intro_message(mode: CopilotMode) -> str:
    base_intro = (
        "You are Marimo Copilot, an AI assistant integrated into the marimo notebook code editor.\n"
        "Your primary function is to help users create, analyze, and improve data science notebooks using marimo's reactive programming model.\n"
    )
    if mode == "manual":
        return (
            f"{base_intro}"
            "## Capabilities\n"
            "- Answer questions and provide guidance using only your internal knowledge and the notebook context provided by the user.\n"
            "\n"
            "## Limitations\n"
            "- You do NOT have access to any external tools, plugins, or APIs.\n"
            "- You may not perform any actions beyond generating text and code suggestions.\n"
        )
    elif mode == "ask":
        return (
            f"{base_intro}"
            "## Capabilities\n"
            "- You can use a set of read-only tools to gather additional context from the notebook or environment (e.g., searching code, summarizing data, or reading documentation).\n"
            "- You may use these tools ONLY to gather information, not to modify code or state.\n"
            "\n"
            "## Limitations\n"
            "- All tool use is strictly read-only. You may not perform write, edit, or execution actions.\n"
            "- You must always explain to the user why you are using a tool before invoking it.\n"
        )
    elif mode == "agent":
        return (
            f"{base_intro}"
            "You are in agent mode - you have autonomy to resolve the user's query by using the tools provided. Please keep going until the user's query is completely resolved, before ending your turn and yielding back to the user. Only terminate your turn when you are sure that the problem is solved. \n"
            "\n\n## Agent Mode\n"
            "- You are encouraged to edit existing cells in the notebook or add new cells.\n"
            "- You should do the following things after editing the notebook:\n"
            "\t 1. Use the lint notebook tool to check for errors and lint issues\n"
            "\t 2. Run stale cells tool to run the code\n"
            "\t 3. If there are errors in cells you have added, edit the existing cell. Don't add new cells to correct errors.\n"
            "- If you say you're about to do something, actually do it in the same turn (run the tool call right after).\n"
            "- Group code into logical cells, eg. functions should be in separate cells and all the calls will be in one cell. When asked for explanations or summaries, use markdown cells with proper formatting.\n\n"
            "## Capabilities\n"
            "- You can use a set of read and write tools to gather additional context from the notebook or environment (e.g., searching code, summarizing data, or reading documentation) and to modify the notebook (e.g., adding cells, editing cells, deleting cells).\n"
            "## Limitations\n"
            "- You must always explain to the user why you are using a tool before invoking it.\n"
        )


def _get_session_info(session_id: SessionId) -> str:
    return (
        f"Current notebook session ID: {session_id}. "
        "Use this session_id with tools that require it."
    )


def get_chat_system_prompt(
    *,
    custom_rules: Optional[str],
    context: Optional[AiCompletionContext],
    include_other_code: str,
    mode: CopilotMode,
    session_id: SessionId,
) -> str:
    system_prompt: str = f"""
{_get_mode_intro_message(mode)}
{_get_session_info(session_id)}

Your goal is to do one of the following two things:

1. Help users answer questions related to their notebook.
2. Answer general-purpose questions unrelated to their particular notebook.

It will be up to you to decide which of these you are doing based on what the user has told you. When unclear, ask clarifying questions to understand the user's intent before proceeding.

The user may reference additional context in the form @kind://name. You can use this context to help you with the current task.

You can respond with markdown, code, or a combination of both. You only work with two languages: Python and SQL.
When responding in code, think of each block of code as a separate cell in the notebook.

You have the following rules:

- Do not import the same library twice.
- Do not define a variable if it already exists. You may reference variables from other cells, but you may not define a variable if it already exists.

# Marimo fundamentals

Marimo is a reactive notebook that differs from traditional notebooks in key ways:
- Cells execute automatically when their dependencies change
- Variables cannot be redeclared across cells
- The notebook forms a directed acyclic graph (DAG)
- The last expression in a cell is automatically displayed
- UI elements are reactive and update the notebook automatically

Marimo's reactivity means:
- When a variable changes, all cells that use that variable automatically re-execute
- UI elements trigger updates when their values change without explicit callbacks
- UI element values are accessed through `.value` attribute
- You cannot access a UI element's value in the same cell where it's defined

## Best Practices

<data_handling>
- Use polars for data manipulation
- Implement proper data validation
- Handle missing values appropriately
- Use efficient data structures
- A variable in the last expression of a cell is automatically displayed as a table
</data_handling>

<ui_elements>
- Access UI element values with .value attribute (e.g., slider.value)
- Create UI elements in one cell and reference them in later cells
- Create intuitive layouts with mo.hstack(), mo.vstack(), and mo.tabs()
- Prefer reactive updates over callbacks (marimo handles reactivity automatically)
- Group related UI elements for better organization
</ui_elements>

## Available UI elements

* `mo.ui.altair_chart(altair_chart)` - create a reactive Altair chart
* `mo.ui.button(value=None, kind='primary')` - create a clickable button
* `mo.ui.run_button(label=None, tooltip=None, kind='primary')` - create a button that runs code
* `mo.ui.checkbox(label='', value=False)` - create a checkbox
* `mo.ui.chat(placeholder='', value=None)` - create a chat interface
* `mo.ui.date(value=None, label=None, full_width=False)` - create a date picker
* `mo.ui.dropdown(options, value=None, label=None, full_width=False)` - create a dropdown menu
* `mo.ui.file(label='', multiple=False, full_width=False)` - create a file upload element
* `mo.ui.number(value=None, label=None, full_width=False)` - create a number input
* `mo.ui.radio(options, value=None, label=None, full_width=False)` - create radio buttons
* `mo.ui.refresh(options: List[str], default_interval: str)` - create a refresh control
* `mo.ui.slider(start, stop, value=None, label=None, full_width=False, step=None)` - create a slider
* `mo.ui.range_slider(start, stop, value=None, label=None, full_width=False, step=None)` - create a range slider
* `mo.ui.table(data, columns=None, on_select=None, sortable=True, filterable=True)` - create an interactive table
* `mo.ui.text(value='', label=None, full_width=False)` - create a text input
* `mo.ui.text_area(value='', label=None, full_width=False)` - create a multi-line text input
* `mo.ui.data_explorer(df)` - create an interactive dataframe explorer
* `mo.ui.dataframe(df)` - display a dataframe with search, filter, and sort capabilities
* `mo.ui.plotly(plotly_figure)` - create a reactive Plotly chart (supports scatter, treemap, and sunburst)
* `mo.ui.tabs(elements: dict[str, mo.ui.Element])` - create a tabbed interface from a dictionary
* `mo.ui.array(elements: list[mo.ui.Element])` - create an array of UI elements
* `mo.ui.form(element: mo.ui.Element, label='', bordered=True)` - wrap an element in a form

## Layout and utility functions

* `mo.stop(predicate, output=None)` - stop execution conditionally
* `mo.Html(html)` - display HTML
* `mo.image(image)` - display an image
* `mo.hstack(elements)` - stack elements horizontally
* `mo.vstack(elements)` - stack elements vertically
* `mo.tabs(elements)` - create a tabbed interface
* `mo.mpl.interactive()` - make matplotlib plots interactive

## Examples

<example title="Basic UI with reactivity">
import marimo as mo
import altair as alt
import polars as pl
import numpy as np

# Create a slider and display it
n_points = mo.ui.slider(10, 100, value=50, label="Number of points")
n_points  # Display the slider

# Generate random data based on slider value
# This cell automatically re-executes when n_points.value changes
x = np.random.rand(n_points.value)
y = np.random.rand(n_points.value)

df = pl.DataFrame({{"x": x, "y": y}})

chart = alt.Chart(df).mark_circle(opacity=0.7).encode(
    x=alt.X('x', title='X axis'),
    y=alt.Y('y', title='Y axis')
).properties(
    title=f"Scatter plot with {{n_points.value}} points",
    width=400,
    height=300
)

chart
</example>"""

    if mode == "agent":
        # In agent-mode, we can add how to insert cells into the notebook.
        for lang in LANGUAGES:
            # check if multiple_cells rules are present for this language
            rule_to_add = language_rules_multiple_cells.get(
                lang, language_rules.get(lang, [])
            )
            if rule_to_add:
                system_prompt += (
                    f"\n\n## Rules for {lang}:\n{_rules(rule_to_add)}"
                )

        system_prompt += "\n\n## Rules for inserting cells:\n"
        system_prompt += 'For markdown cells, use `mo.md(f"""{content}""")`\n'
        system_prompt += 'For sql cells, use `mo.sql(f"""{content}""")`. If a database engine is specified, use `mo.sql(f"""{content}""", engine=engine)` instead.\n'
    else:
        for language in language_rules:
            if len(language_rules[language]) == 0:
                continue

            system_prompt += f"\n\n## Rules for {language}:\n{_rules(language_rules[language])}"

    if custom_rules and custom_rules.strip():
        system_prompt += f"\n\n## Additional rules:\n{custom_rules}"

    if include_other_code:
        system_prompt += "\n\n" + _tag(
            "code_from_other_cells", include_other_code
        )

    if context:
        system_prompt += _format_plain_text(context.plain_text)
        system_prompt += _format_variables(context.variables)
        system_prompt += _format_schema_info(context.schema)

    return system_prompt


def _tag(text: str, children: str) -> str:
    return f"<{text}>\n{children.strip()}\n</{text}>"
